import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from tqdm import tqdm

class ANFISLayer(nn.Module):
    def __init__(self, input_dim, S=2, T=16, learning_rate=0.01):
        super(ANFISLayer, self).__init__()
        self.S = S
        self.T = T
        self.learning_rate = learning_rate

        # Initialize parameters
        self.c1 = nn.Parameter(torch.empty(S, input_dim).uniform_(-1, 1))
        self.b1 = nn.Parameter(torch.empty(S, input_dim).uniform_(-1, 1))
        self.w1 = nn.Parameter(torch.empty(1, T).uniform_(0, 1))

    def forward(self, samplein):
        # Normalization
        samplein_min = samplein.min(dim=1, keepdim=True).values
        samplein_max = samplein.max(dim=1, keepdim=True).values
        sampleinnorm = (samplein - samplein_min) / (samplein_max - samplein_min)

        # Initialize output tensor
        Y = torch.zeros_like(sampleinnorm)

        u1 = torch.zeros((self.S, sampleinnorm.shape[1]), device=sampleinnorm.device)
        alpha = []

        # Compute membership degrees and output
        for m in range(sampleinnorm.shape[0]):
            for i in range(self.S):
                for j in range(sampleinnorm.shape[1]):
                    u1[i][j] = torch.exp(-((sampleinnorm[m, j] - self.c1[i, j]) ** 2) / (self.b1[i, j] ** 2))

            # Rule generation layer
            alpha1 = torch.zeros(self.T, device=sampleinnorm.device)
            for i in range(self.S):
                a = 1 if i != 1 else 8
                for p in range(self.S):
                    b = 1 if p != 1 else 4
                    for q in range(self.S):
                        c = 1 if q != 1 else 2
                        for k in range(self.S):
                            idx = i * a + p * b + q * c + k
                            alpha1[idx] = u1[i][0] * u1[p][1] * u1[q][2] * u1[k][3]

            alpha.append(alpha1)
            alphasum = alpha1.sum()  # Scalar value
            # Output layer calculation
            Y[m] = torch.dot(self.w1.squeeze(), alpha1) / alphasum if alphasum != 0 else 0  # Avoid division by zero

        return Y


class FCNN(nn.Module):
    def __init__(self, n_femap=4, stride=2, mu=3.0, sigma=1.2, dropout=True, num_classes=3):
        super(FCNN, self).__init__()
        self.dropout = dropout
        self.conv1 = nn.Conv2d(3, 20, kernel_size=5, stride=2, padding=1)
        self.conv2 = nn.Conv2d(20, 40, kernel_size=3, stride=2, padding=1)
        self.conv3 = nn.Conv2d(40, 40, kernel_size=3, stride=2, padding=1)
        self.conv4 = nn.Conv2d(40, 40, kernel_size=3, stride=2, padding=1)
        self.conv5 = nn.Conv2d(40, n_femap, kernel_size=5, stride=stride, padding=1)
        self.anfis_layer = ANFISLayer(64)  # Use the ANFIS layer initialized here
        if self.dropout:
            self.dropout_layer = nn.Dropout(0.2)

        self.fc1 = nn.Linear(24*64, 128)  # Update based on the feature map size
        self.linear1 = nn.Linear(128, num_classes)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, kernel_size=3, stride=1, padding=1)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, kernel_size=2, stride=1, padding=1)
        x = F.relu(self.conv3(x))
        x = F.max_pool2d(x, kernel_size=2, stride=1, padding=1)
        x = F.relu(self.conv4(x))
        x = F.max_pool2d(x, kernel_size=2, stride=1, padding=1)
        x = F.relu(self.conv5(x))


        if self.dropout:
            x = self.dropout_layer(x)

        x = x.view(-1, 8 * 8)
        print(x.shape)
        x = self.anfis_layer(x)  # 使用 self.anfis_layer，传入相应的输入和输出
        x = x.view(-1)
        x =  self.fc1(x)
        out = self.linear1(x)

        return out


# Example usage
if __name__ == "__main__":
    # Create the model
    model = FCNN(n_femap=4, stride=2, mu=3.0, sigma=1.2, dropout=True)

    # Create a sample input (batch size of 2, 3 channels, 32x32 image)
    sample_input = torch.randn(6, 3, 256, 256)

    # Get model output
    output = model(sample_input)
    print("Output shape:", output.shape)  # Should be (2, num_classes)
